# encoding: utf-8

# Author     : forward4<liang.ma@siat.ac.cn>
# Datetime   :
# Project    : monkey_neuralDecoding
# File       : utils.py
# description: 文件说明

import os
import random
import shutil

import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import find_peaks, butter, sosfiltfilt

from color_map import cnames


def create_signal(sample_rate,freqs):
    t = np.linspace(0, 1, sample_rate)

    signals = [np.sin(2 * np.pi * f * t) for f in range(freqs)]
    combined_signal = sum(signals)
    return combined_signal


def reset_dir(raw_dir):
    shutil.rmtree(raw_dir)


def new_dir(base_dir, new_dir_name):
    """
    Recreate a new directory, even though it already exists.
    :param base_dir:
    :param new_dir_name:
    :return:
    """
    new_dir = os.path.join(base_dir, new_dir_name)
    if os.path.exists(new_dir):
        reset_dir(new_dir)
    os.mkdir(new_dir)
    return new_dir

def find_spikes(data, threshold=60, fs=30000):
    peaks, _ = find_peaks(-data, height=threshold, distance=int(fs * 0.001))
    return peaks

def lfp_bandpass(lfp, low=250, high=7000, fs=20000, order=6):
    w0 = 2 * low / fs
    w1 = 2 * high / fs
    sos = butter(order, [w0, w1], btype='band', output='sos')
    filtered = sosfiltfilt(sos, lfp)
    return filtered


def get_all_dirs(base_dir=None, suffix="rhd", prefix=""):
    """
    Get all existing directories with specific suffix and prefix in current base directory
    :param base_dir:
    :param suffix:
    :param prefix:
    :return:
    """
    if not base_dir:
        base_dir = "../data"
    file_name = []
    file_dir = []
    for roots, dirs, files in os.walk(base_dir, topdown=False, onerror=None, followlinks=False):
        if len(files):
            for file in files:
                if suffix and file.endswith(suffix):
                    file_name.append(os.path.join(roots, file))
                    file_dir.append(roots)
                if prefix and file.startswith(prefix):
                    file_name.append(os.path.join(roots, file))
                    file_dir.append(roots)
        else:
            pass
    return file_name, file_dir


def view_trace_with_marker(data, time, marker=None, fs=30000, unit="n"):
    """
    :param data: a numpy array with shape of (1, n_points) or (n_points, )
    :param time: a tuple or list with size of 2 elements. The first element denotes the startup time,
    and the second element denotes the ending time.
    :param marker:
    :return: None
    """
    size_key = {
        'xy_fontsize': 25,
        'tick_size': 20,
    }
    # n_point = np.ceil(time[-1])
    time_m = np.ceil(time[-1])
    # step = 100 if unit == "s" else 100 * fs
    # timex = np.arange(0, time_m, step)
    timex = np.linspace(0, time_m, 5)
    timex_label = [str(np.ceil(i / fs)) for i in timex]
    ymin = np.min(data)
    ymax = np.max(data)

    fig = plt.figure(figsize=[16, 6])
    ax = plt.subplot()
    # ax.plot(time, data, c='k', linewidth=1)
    ax.plot(data, c='k', linewidth=1)
    if marker and isinstance(marker, (list, np.ndarray)) and len(marker) > 0:
        ax.vlines(time[marker], ymin=ymin, ymax=ymax, colors='red', linestyles='--',
                  linewidth=2)  # -1 for amending beginning
    ax.set_title('Spike Trace', fontsize=size_key['xy_fontsize'], fontweight='bold')
    ax.set_xlim(0, time_m)
    ax.set_xticks(timex)
    ax.set_xticklabels(timex_label)

    ax.set_xlabel(f'Time (s)', fontsize=size_key['xy_fontsize'], fontweight='bold')
    ax.set_ylabel(f'Amplitute ({chr(956)}V)', fontsize=size_key['xy_fontsize'], fontweight='bold')
    ax.tick_params(axis='both', labelsize=size_key['tick_size'], length=6, width=3)
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    ax.spines['bottom'].set_linewidth(4)
    ax.spines['left'].set_linewidth(4)


def choose_color_random():
    cn_list = list(cnames.keys())
    cn_num = len(cn_list)
    ci = random.randint(0, cn_num - 1)
    return cn_list[ci]


def plot_spike(spikes, linewidth=0.5, layer=1):
    """
    :param spikes:
    :param linewidth:
    :param layer:
    :return:
    """
    # plt.figure()
    ax = plt.subplot()
    spike_num = spikes.shape[0]
    for sid in range(spike_num):
        spike = spikes[sid]
        ax.plot(spike, c=choose_color_random(), linewidth=linewidth, zorder=layer)
    # if is_save:
    #     save_dir = os.path.join(root_dir, "spike_img")
    #     if not os.path.exists(save_dir):
    #         os.mkdir(save_dir)
    #     plt.savefig(os.path.join(save_dir, f"{file_name}.png"))
    #     plt.close()
    # else:
    #     plt.show()

if __name__ == '__main__':
    a = create_signal(1000,500)
